package core

import (
	"github.com/nemars/heapmap"
	"sync"
)

type SpanRecord struct {
	SpanID       int64
	PktRecordMap map[string]*PktRecord // key is key of PktRecord
}

type PktRecord struct {
	Key    string
	SpanID int64

	SrcIP   string
	DstIP   string
	SrcPort uint16
	DstPort uint16
	Proto   string
	Tos     uint8

	UseLowLevelTxTs bool
	TxTs            int64
	RxTs            int64
	RxTxTs          int64
	Latency         int64
	Received        bool
	// Code StatusCode
	BitFlip bool
}

type AggreRecord struct {
	Key string

	SrcIP    string
	DstIP    string
	SrcPort  int
	DstPort  int
	Protocol string
	Tos      int

	TxPktCnt   int
	LossPktCnt int
	BitFlipCnt int
	MinLatency int64
	AvgLatency int64
	MaxLatency int64

	PktRecords []*PktRecord
}

type AggreSpanRecord struct {
	SpanID         int64
	AggreRecordMap map[string]*AggreRecord // key is target
}

func NewSpanRecord(spanID int64) *SpanRecord {
	return &SpanRecord{
		SpanID:       spanID,
		PktRecordMap: make(map[string]*PktRecord),
	}
}

func NewAggreSpanRecord(spanID int64) *AggreSpanRecord {
	return &AggreSpanRecord{
		SpanID:         spanID,
		AggreRecordMap: make(map[string]*AggreRecord),
	}
}

type PktRecordBucket struct {
	M       sync.RWMutex
	HeapMap heapmap.HeapMap[string, *PktRecord, int64]
}

func NewPktRecordBucket() *PktRecordBucket {
	return &PktRecordBucket{
		HeapMap: heapmap.NewMin[string, *PktRecord, int64](),
	}
}

func (b *PktRecordBucket) Set(v *PktRecord) {
	b.M.Lock()
	defer b.M.Unlock()
	b.HeapMap.Set(v.Key, v, v.SpanID)
}

func (b *PktRecordBucket) Update(v *PktRecord) {
	b.M.RLock()
	entry, ok := b.HeapMap.Get(v.Key)
	b.M.RUnlock()
	if ok {
		entry.Value.Received = v.Received
		entry.Value.BitFlip = v.BitFlip
		entry.Value.RxTs = v.RxTs
		entry.Value.Latency = v.RxTs - entry.Value.TxTs
	}
}

func (b *PktRecordBucket) Peek() *PktRecord {
	b.M.RLock()
	entry, ok := b.HeapMap.Peek()
	b.M.RUnlock()
	if ok {
		return entry.Value
	}
	return nil
}

func (b *PktRecordBucket) Pop() *PktRecord {
	b.M.Lock()
	entry, ok := b.HeapMap.Pop()
	b.M.Unlock()
	if ok {
		return entry.Value
	}
	return nil
}
